use memmap2::{MmapMut, MmapOptions};
use std::fs::{File, OpenOptions};
use std::io::{Result, Write, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock};

const CHUNK_SIZE: usize = 4 * 1024 * 1024; // 4MB
const NUM_CHUNKS: usize = 256; // 256 chunks = 1GB

/// 管理延迟映射的内存区域
struct LazyMappedMemory {
    base_path: PathBuf,
    mappings: Vec<Option<MmapMut>>,
    file_handles: Vec<Option<File>>,
}

impl LazyMappedMemory {
    fn new<P: AsRef<Path>>(base_path: P) -> Result<Self> {
        // 确保基础目录存在
        std::fs::create_dir_all(&base_path)?;
        
        // 使用迭代器初始化 mappings 和 file_handles
        let mappings = (0..NUM_CHUNKS).map(|_| None).collect();
        let file_handles = (0..NUM_CHUNKS).map(|_| None).collect();
        
        Ok(Self {
            base_path: base_path.as_ref().to_path_buf(),
            mappings,
            file_handles,
        })
    }
    
    // 其他方法保持不变...
    fn get_byte(&mut self, offset: usize) -> Result<u8> {
        if offset >= CHUNK_SIZE * NUM_CHUNKS {
            return Ok(0);
        }
        
        let chunk_idx = offset / CHUNK_SIZE;
        let chunk_offset = offset % CHUNK_SIZE;
        
        // 如果映射不存在，尝试创建或返回0
        if self.mappings[chunk_idx].is_none() {
            if let Ok(mapping) = self.ensure_mapping(chunk_idx, false) {
                self.mappings[chunk_idx] = Some(mapping);
                return Ok(self.mappings[chunk_idx].as_ref().unwrap()[chunk_offset]);
            } else {
                // 文件不存在，返回0
                return Ok(0);
            }
        }
        
        Ok(self.mappings[chunk_idx].as_ref().unwrap()[chunk_offset])
    }
    
    fn set_byte(&mut self, offset: usize, value: u8) -> Result<()> {
        if offset >= CHUNK_SIZE * NUM_CHUNKS {
            return Ok(());
        }
        
        let chunk_idx = offset / CHUNK_SIZE;
        let chunk_offset = offset % CHUNK_SIZE;
        
        // 确保映射存在（如果需要则创建文件）
        if self.mappings[chunk_idx].is_none() {
            let mapping = self.ensure_mapping(chunk_idx, true)?;
            self.mappings[chunk_idx] = Some(mapping);
        }
        
        // 写入数据
        self.mappings[chunk_idx].as_mut().unwrap()[chunk_offset] = value;
        Ok(())
    }
	    /// 将所有修改刷新到磁盘
    fn flush_all(&mut self) -> Result<()> {
        for i in 0..NUM_CHUNKS {
            if let Some(mapping) = self.mappings[i].as_mut() {
                mapping.flush()?;
            }
        }
        Ok(())
    }
    
    /// 将指定块刷新到磁盘
    fn flush_chunk(&mut self, chunk_idx: usize) -> Result<()> {
        if let Some(mapping) = self.mappings[chunk_idx].as_mut() {
            mapping.flush()?;
        }
        Ok(())
    }
    
    /// 卸载指定块（释放内存映射）
    fn unload_chunk(&mut self, chunk_idx: usize) {
        self.mappings[chunk_idx] = None;
        // 注意：我们保留文件句柄以便后续快速重新映射
    }
	
    fn ensure_mapping(&mut self, chunk_idx: usize, create_if_missing: bool) -> Result<MmapMut> {
        let file_path = self.base_path.join(format!("chunk_{:03}.bin", chunk_idx));
        
        // 检查文件是否存在
        if !file_path.exists() {
            if !create_if_missing {
                // 文件不存在且不允许创建，返回错误
                return Err(std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    "File does not exist",
                ));
            }
            
            // 创建文件并初始化为全零
            let mut file = OpenOptions::new()
                .read(true)
                .write(true)
                .create(true)
                .open(&file_path)?;
            
            // 设置文件大小为 CHUNK_SIZE
            file.set_len(CHUNK_SIZE as u64)?;
            
            // 初始化文件内容为全零（可选，但可以确保一致性）
            let zeros = vec![0u8; CHUNK_SIZE];
            file.seek(SeekFrom::Start(0))?;
            file.write_all(&zeros)?;
            file.flush()?;
            
            self.file_handles[chunk_idx] = Some(file);
        } else if self.file_handles[chunk_idx].is_none() {
            // 文件存在但尚未打开
            let file = OpenOptions::new()
                .read(true)
                .write(true)
                .open(&file_path)?;
            self.file_handles[chunk_idx] = Some(file);
        }
        
        // 创建内存映射
        unsafe {
            let mapping = MmapOptions::new()
                .map_mut(self.file_handles[chunk_idx].as_ref().unwrap())?;
            Ok(mapping)
        }
    }
    
    // 其他方法保持不变...
}

// 提供线程安全版本的包装器
#[derive(Clone)]
struct ThreadSafeLazyMappedMemory {
    inner: Arc<RwLock<LazyMappedMemory>>,
}

impl ThreadSafeLazyMappedMemory {
    fn new<P: AsRef<Path>>(base_path: P) -> Result<Self> {
        Ok(Self {
            inner: Arc::new(RwLock::new(LazyMappedMemory::new(base_path)?)),
        })
    }
    
    fn get_byte(&self, offset: usize) -> Result<u8> {
        self.inner.write().unwrap().get_byte(offset)
    }
    
    fn set_byte(&self, offset: usize, value: u8) -> Result<()> {
        self.inner.write().unwrap().set_byte(offset, value)
    }
    
    fn flush_all(&self) -> Result<()> {
        self.inner.write().unwrap().flush_all()
    }
}

fn main() -> Result<()> {
    let base_path = "./lazy_mapped_data";
    
    // 创建延迟映射的内存管理器
    let memory = ThreadSafeLazyMappedMemory::new(base_path)?;
    
    // 示例：写入一些数据（会自动创建文件）
    for i in 0..CHUNK_SIZE*5-1 {
        memory.set_byte(i, (i % 256) as u8)?;
    }
    
    // 示例：读取数据（未创建的文件区域返回0）
    for i in (CHUNK_SIZE * 4 - 10) .. (CHUNK_SIZE * 4 + 10) {
        let value = memory.get_byte(i)?;
        println!("Offset {}: {}", i, value);
    }
    
    // 示例：读取未创建的文件区域（返回0）
    let unallocated_value = memory.get_byte(CHUNK_SIZE * 200)?;
    println!("Unallocated area value: {}", unallocated_value);
    
    // 刷新所有更改到磁盘
    memory.flush_all()?;
    
    println!("延迟映射设置完成");
    
    Ok(())
}

/* cargo watch -c -q -w ./src -x "run" */
